﻿using System;
using System.Collections.Generic;
using System.IdentityModel.Tokens;
using System.IO;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using System.Web;
using System.Xml;
using Microsoft.IdentityModel.SecurityTokenService;
using Microsoft.IdentityModel.Tokens.Saml11;
using Tokens;
using WifTokenTypes = Microsoft.IdentityModel.Tokens.SecurityTokenTypes;

namespace TokenUtil
{
   class Program
   {
      static Parameters _parameters = new Parameters(
         new Dictionary<string, Parameter>() {
            { ParameterFlags.ADFS, new Parameter() { Type=null, Description="Request token from ADFS STS" } },
            { ParameterFlags.ACS, new Parameter() { Type=null, Description="Request token from ACS STS" } },
            { ParameterFlags.TYPE, new Parameter() { Type=typeof(string), Description="Token type", Value="saml" } },
            { ParameterFlags.HOSTNAME, new Parameter() { Type=typeof(string), Description="STS hostname" } },
            { ParameterFlags.APPLIESTO, new Parameter() { Type=typeof(string), Description="Realm that the token applies to" } },
            { ParameterFlags.USERNAME, new Parameter() { Type=typeof(string), Description="Username" } },
            { ParameterFlags.PASSWORD, new Parameter() { Type=typeof(string), Description="Password", Value=string.Empty } },
            { ParameterFlags.PFX, new Parameter() { Type=typeof(string), Description="Path to decryption pfx" } },
            { ParameterFlags.PFX_PASSWORD, new Parameter() { Type=typeof(string), Description="Path to decryption pfx password" } },
            { ParameterFlags.TOKEN, new Parameter() { Type=typeof(string), Description="A token" } },
            { ParameterFlags.KEYTYPE, new Parameter() { Type=typeof(string), Description="Token key type" } } },
         new List<string>() { ParameterFlags.ADFS, ParameterFlags.ACS },
         new List<Action<Dictionary<string, Parameter>>>() {
            { ps => {
               if (!ps[ParameterFlags.ADFS].IsSet && !ps[ParameterFlags.ACS].IsSet)
                  throw new InvalidOperationException(
                     "You must specify either the /" + ParameterFlags.ADFS + " or /" + ParameterFlags.ACS + " argument.");

               if (!ps[ParameterFlags.HOSTNAME].IsSet)
                  throw new InvalidOperationException(
                     "You must include the target STS hostname with the /" + ParameterFlags.HOSTNAME + " argument.");

               if (ps[ParameterFlags.ADFS].IsSet && ps[ParameterFlags.TYPE].IsSet)
                  throw new InvalidOperationException(
                     "The argument " + ParameterFlags.TYPE + " is not supported by for token requests to ADFS.");

               if (ps[ParameterFlags.ACS].IsSet)
               {
                  if (ps[ParameterFlags.TYPE].Value.Equals("swt") && String.IsNullOrEmpty((string)ps[ParameterFlags.USERNAME].Value))
                     throw new InvalidOperationException(
                        "Requesting an SWT token from ACS requires a username and password.");

                  if (!ps[ParameterFlags.TYPE].Value.Equals("swt") && !ps[ParameterFlags.TYPE].Value.Equals("saml"))
                     throw new InvalidOperationException(
                        ps[ParameterFlags.TYPE].Value + " is not a recognized token type. Please specify either SAML or SWT.");
               }
            } } } );

      static SecurityToken _lastToken;
      static string _lastKeyType;

      static void Main(string[] args)
      {
         try
         {
            _parameters.Process(args);

            foreach (Dictionary<string, Parameter> parameterSet in _parameters.ParameterSets)
            {
               if (parameterSet[ParameterFlags.ADFS].IsSet)
               {
                  string keyType;
                  if (parameterSet[ParameterFlags.KEYTYPE].IsSet)
                     keyType = ResolveKeyType(parameterSet[ParameterFlags.KEYTYPE].Value as string);
                  else
                     keyType = KeyTypes.Symmetric;

                  var token = AdfsStsClient.GetTokenWsTrust13(
                     parameterSet[ParameterFlags.HOSTNAME].Value as string,
                     parameterSet[ParameterFlags.APPLIESTO].Value as string,
                     keyType,
                     parameterSet[ParameterFlags.USERNAME].Value as string,
                     parameterSet[ParameterFlags.PASSWORD].Value as string);

                  //var token = AdfsStsClient.GetTokenWsTrust2005(
                  //   parameterSet[ParameterFlags.HOSTNAME].Value as string,
                  //   parameterSet[ParameterFlags.APPLIESTO].Value as string,
                  //   keyType,
                  //   parameterSet[ParameterFlags.USERNAME].Value as string,
                  //   parameterSet[ParameterFlags.PASSWORD].Value as string);

                  var xmlToken = token as GenericXmlSecurityToken;
                  _lastToken = xmlToken;
                  _lastKeyType = keyType;

                  if (!parameterSet[ParameterFlags.PFX].IsSet)
                  {
                     Console.Write(xmlToken.TokenXml.OuterXml);
                  }
                  else
                  {
                     //Decrypting the token requires access to the private key of the certificate which contains the public key used to encrypt it.
                     //Private keys are typically distributed with their corresponding cerficates in password protected PFX files.
                     var pfx = new X509Certificate2(
                        parameterSet[ParameterFlags.PFX].Value as string,
                        parameterSet[ParameterFlags.PFX_PASSWORD].Value as string);

                     var tokenXmlDoc = new XmlDocument();
                     {
                        using (var stringWriter = new StringWriter())
                        {
                           // You also need a XmlTextWriter to work with the StringWriter
                           using (var xmlWriter = new XmlTextWriter(stringWriter))
                           {
                              // Write the XML to the XMLTextWriter
                              xmlToken.TokenXml.WriteTo(xmlWriter);
                              // Use the StringWriter to make it a string
                              var temp = stringWriter.ToString();
                              // Then Load the XML into the XmlDocument
                              tokenXmlDoc.LoadXml(temp);
                           }
                        }
                     }

                     var nsmgr = new XmlNamespaceManager(tokenXmlDoc.NameTable);
                     nsmgr.AddNamespace("xenc", "http://www.w3.org/2001/04/xmlenc#");

                     var nodeList = tokenXmlDoc.SelectNodes("descendant::xenc:CipherValue", nsmgr);
                     var xmlElSymmetricKeyCipher = (XmlElement)nodeList[0];

                     var base64SymmetricKeyCipher = xmlElSymmetricKeyCipher.InnerText;
                     var symmetricKeyCipherBytes = Convert.FromBase64String(base64SymmetricKeyCipher);

                     using (var rsa = (RSACryptoServiceProvider)pfx.PrivateKey)
                     {
                        var symmetricKey = rsa.Decrypt(symmetricKeyCipherBytes, true);

                        using (var alg = Rijndael.Create())
                        {
                           alg.Key = symmetricKey;

                           var xmlElTokenCipher = (XmlElement)nodeList[1];
                           string base64TokenCipher = xmlElTokenCipher.InnerText;
                           byte[] tokenRawCipher = Convert.FromBase64String(base64TokenCipher);

                           // No using necessary for this MemoryStream because the CryptoStream using will dispose it as well.
                           var memoryStream = new MemoryStream();
                           using (var cryptoStream = new CryptoStream(memoryStream, alg.CreateDecryptor(), CryptoStreamMode.Write))
                           {
                              cryptoStream.Write(tokenRawCipher, 0, tokenRawCipher.Length);
                              cryptoStream.Flush();

                              var encoder = new UTF8Encoding();
                              string plainText = encoder.GetString(memoryStream.ToArray());

                              // Trim the gobledygook from the beginning.
                              int firstCaret = plainText.IndexOf('<');
                              plainText = plainText.Substring(firstCaret, plainText.Length - firstCaret);
                              plainText = plainText + ":Assertion>";

                              Console.Write(plainText);
                           }
                        }
                     }
                  }
               }
               else if (parameterSet[ParameterFlags.ACS].IsSet)
               {
                  if (parameterSet[ParameterFlags.USERNAME].IsSet && !parameterSet[ParameterFlags.PFX].IsSet)
                  {
                     if (parameterSet[ParameterFlags.TYPE].Value.Equals("saml"))
                     {
                        string keyType;
                        if (parameterSet[ParameterFlags.KEYTYPE].IsSet)
                           keyType = parameterSet[ParameterFlags.KEYTYPE].Value as string;
                        else
                           keyType = KeyTypes.Bearer;

                        SecurityToken token2 = AcsStsClient.GetToken(
                              parameterSet[ParameterFlags.HOSTNAME].Value as string,
                              parameterSet[ParameterFlags.APPLIESTO].Value as string,
                              keyType,
                              parameterSet[ParameterFlags.USERNAME].Value as string,
                              parameterSet[ParameterFlags.PASSWORD].Value as string);

                        _lastToken = token2;
                        _lastKeyType = keyType;
                        GenericXmlSecurityToken xmlToken2 = token2 as GenericXmlSecurityToken;
                        Console.Write(xmlToken2.TokenXml.InnerXml);
                     }
                     else if (parameterSet[ParameterFlags.TYPE].Value.Equals("swt"))
                     {
                        string swtToken = AcsStsClient.GetSwtToken(
                           parameterSet[ParameterFlags.HOSTNAME].Value as string,
                           parameterSet[ParameterFlags.APPLIESTO].Value as string,
                           parameterSet[ParameterFlags.USERNAME].Value as string,
                           parameterSet[ParameterFlags.PASSWORD].Value as string);

                        Console.Write(HttpUtility.UrlDecode(swtToken));
                     }
                     else
                        throw new InvalidOperationException("Unsupported token type.");
                  }
                  else if (parameterSet[ParameterFlags.PFX].IsSet)
                  {
                     byte[] clientCertificateBytes;
                     using (FileStream stream = File.OpenRead(parameterSet[ParameterFlags.PFX].Value as string))
                     {
                        using (BinaryReader reader = new BinaryReader(stream))
                        {
                           clientCertificateBytes = reader.ReadBytes((int)stream.Length);
                        }
                     }

                     string swtToken = AcsStsClient.GetSwtToken2(
                        parameterSet[ParameterFlags.HOSTNAME].Value as string,
                        WifTokenTypes.Saml11TokenProfile11,
                        parameterSet[ParameterFlags.APPLIESTO].Value as string,
                        parameterSet[ParameterFlags.USERNAME].Value as string,
                        clientCertificateBytes,
                        parameterSet[ParameterFlags.PFX_PASSWORD].Value as string);

                     //string swtToken = AcsStsClient.GetSwtToken2b(
                     //   parameterSet[ParameterFlags.HOSTNAME].Value as string,
                     //   parameterSet[ParameterFlags.APPLIESTO].Value as string,
                     //   parameterSet[ParameterFlags.USERNAME].Value as string,
                     //   clientCertificateBytes,
                     //   parameterSet[ParameterFlags.PFX_PASSWORD].Value as string);

                     Console.Write(HttpUtility.UrlDecode(swtToken));
                  }
                  else
                  {
                     if (null == _lastToken)
                     {
                        string tokenStr;
                        if (parameterSet[ParameterFlags.TOKEN].IsSet)
                        {
                           string tokenPath = parameterSet[ParameterFlags.TOKEN].Value as string;
                           using (TextReader reader = new StreamReader(tokenPath))
                              tokenStr = reader.ReadToEnd();
                        }
                        else
                           tokenStr = Console.ReadLine();

                        Saml11SecurityTokenHandler handler = new Saml11SecurityTokenHandler();
                        //Saml2SecurityTokenHandler handler = new Saml2SecurityTokenHandler();
                        handler.Configuration = new Microsoft.IdentityModel.Tokens.SecurityTokenHandlerConfiguration();
                        SecurityToken token;
                        using (XmlReader reader = XmlReader.Create(new StringReader(tokenStr)))
                           token = handler.ReadToken(reader);
                        _lastToken = token;
                     }

                     if (parameterSet[ParameterFlags.TYPE].Value.Equals("saml"))
                     {
                        string keyType;
                        if (parameterSet[ParameterFlags.KEYTYPE].IsSet)
                           keyType = ResolveKeyType(parameterSet[ParameterFlags.KEYTYPE].Value as string);
                        else
                           keyType = KeyTypes.Bearer;

                        SecurityToken token2 = null;
                        if (KeyTypes.Bearer == _lastKeyType)
                        {
                           token2 = AcsStsClient.GetTokenWithIssuedTokenBearer(
                                 parameterSet[ParameterFlags.HOSTNAME].Value as string,
                                 parameterSet[ParameterFlags.APPLIESTO].Value as string,
                                 keyType,
                                 _lastToken);
                        }
                        else if (KeyTypes.Symmetric == _lastKeyType)
                        {
                           token2 = AcsStsClient.GetTokenWithIssuedTokenSymmetric(
                                 parameterSet[ParameterFlags.HOSTNAME].Value as string,
                                 parameterSet[ParameterFlags.APPLIESTO].Value as string,
                                 keyType,
                                 _lastToken);
                        }
                        else
                        {
                           throw new InvalidOperationException("key type not supported");
                        }

                        _lastToken = token2;
                        _lastKeyType = keyType;
                        GenericXmlSecurityToken xmlToken2 = token2 as GenericXmlSecurityToken;
                        Console.Write(xmlToken2.TokenXml.InnerXml);
                     }
                  }
               }
            }
         }
         catch (Exception ex)
         {
            Console.WriteLine(ex.Message);

            // These are helpful messages.
            if (ex.ToString().Contains("ID3037:"))
               Console.WriteLine(
                  "NOTE: When receiving this error from ADFS, you may see this message in the server log:\n" +
                  "\"ID4007: The symmetric key inside the requested security token must be encrypted\"\n" +
                  "Unfortunately, This is message is often misleading. " +
                  "In many cases, the problem is that the configuration of the relying party does " +
                  "not specify an encryption certificate. If you're building the relying party " +
                  "configuration from a FederationMetadata.xml file, you will have to configure " +
                  "your relying party application to specify a token decryption certificate.");

            if (ex.ToString().Contains("ACS10002") && ex.Message.Contains("ACS50008"))
               Console.WriteLine(
                  "NOTE: When receiving this error from ACS the cause is probably that the security token\n" +
                  " signing certificate and encryption certificate don’t match between ADFS and ACS.\n" +
                  "To resolve this you may try updating the federation metadata of ADFS in ACS and\n" +
                  "updating the federation metadata of ACS in ADFS.");
         }
      }

      /// <summary>
      /// Resolve the partial key type name to the full URL value.
      /// </summary>
      /// <param name="keyType">Key type name</param>
      /// <returns>Key type URL</returns>
      static string ResolveKeyType(string keyType)
      {
         switch (keyType)
         {
            case "Bearer":
               return KeyTypes.Bearer;
            case "Symmetric":
               return KeyTypes.Symmetric;
            case "Asymmetric":
               return KeyTypes.Asymmetric;
            default:
               throw new InvalidOperationException("Unrecognized key type.");
         };
      }
   }
}
